package com.cy.ragbase.service;

import com.cy.ragbase.dto.CodeSearchResponse;
import com.cy.ragbase.dto.SearchResult;
import com.cy.ragbase.entity.CodeBlock;
import com.cy.ragbase.entity.DocumentChunk;
import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.dml.DeleteParam;
import io.milvus.response.SearchResultsWrapper;
import io.milvus.grpc.SearchResults;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import io.milvus.param.MetricType;

import java.util.*;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@Slf4j
public class MilvusService {

    private final MilvusServiceClient milvusClient;

    @Value("${milvus.collection-name}")
    private String collectionName;

    /**
     * 批量插入文档向量
     */
    public void insertDocumentChunks(List<DocumentChunk> chunks) {
        if (chunks.isEmpty()) {
            return;
        }

        try {
            List<String> ids = chunks.stream()
                    .map(DocumentChunk::getId)
                    .collect(Collectors.toList());

            List<String> documentIds = chunks.stream()
                    .map(DocumentChunk::getDocumentId)
                    .collect(Collectors.toList());

            List<Long> chunkIndexes = chunks.stream()
                    .map(chunk -> (long) chunk.getChunkIndex())
                    .collect(Collectors.toList());

            List<String> contents = chunks.stream()
                    .map(DocumentChunk::getContent)
                    .collect(Collectors.toList());

            List<List<Float>> vectors = chunks.stream()
                    .map(chunk -> {
                        List<Float> vector = new ArrayList<>();
                        for (float f : chunk.getVector()) {
                            vector.add(f);
                        }
                        return vector;
                    })
                    .collect(Collectors.toList());

            List<InsertParam.Field> fields = Arrays.asList(
                    new InsertParam.Field("id", ids),
                    new InsertParam.Field("document_id", documentIds),
                    new InsertParam.Field("chunk_index", chunkIndexes),
                    new InsertParam.Field("content", contents),
                    new InsertParam.Field("vector", vectors)
            );

            InsertParam insertParam = InsertParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withFields(fields)
                    .build();

            milvusClient.insert(insertParam);
            log.info("向Milvus插入了 {} 个文档块", chunks.size());

        } catch (Exception e) {
            log.error("向Milvus插入文档块失败", e);
            throw new RuntimeException("Failed to insert into Milvus", e);
        }
    }

    /**
     * 搜索相似向量
     */
    public List<SearchResult> searchSimilarChunks(float[] queryVector, int topK, float threshold) {
        try {
            List<String> outputFields = Arrays.asList("id", "document_id", "chunk_index", "content");


            List<Float> vectorList = new ArrayList<>(queryVector.length);
            for (float f : queryVector) {
                vectorList.add(f);
            }
            List<List<Float>> searchVectors = Collections.singletonList(vectorList);

            SearchParam searchParam = SearchParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withMetricType(MetricType.COSINE)
                    .withOutFields(outputFields)
                    .withTopK(topK)
                    .withVectors(searchVectors)
                    .withVectorFieldName("vector")
                    .withParams("{\"nprobe\":10}")
                    .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                    .build();

            SearchResults searchResults = milvusClient.search(searchParam).getData();
            SearchResultsWrapper wrapper = new SearchResultsWrapper(searchResults.getResults());

            List<SearchResult> results = new ArrayList<>();


            if (wrapper.getRowRecords().isEmpty()) {
                log.info("未找到搜索结果");
                return results;
            }


            // 获取第一个查询的结果（我们只发送了一个查询向量）
            List<SearchResultsWrapper.IDScore> idScores = wrapper.getIDScore(0);

            List<?> documentIdList = wrapper.getFieldData("document_id", 0);
            List<?> contentList = wrapper.getFieldData("content", 0);
            List<?> chunkIndexList = wrapper.getFieldData("chunk_index", 0);


            for (int i = 0; i < idScores.size(); i++) {
                SearchResultsWrapper.IDScore idScore = idScores.get(i);
                float score = idScore.getScore();

                if (score >= threshold) {
                    SearchResult result = new SearchResult();
                    result.setScore(score);

                    try {
                        // 安全获取字段数据
                        if (documentIdList != null && i < documentIdList.size()) {
                            Object documentIdObj = documentIdList.get(i);
                            if (documentIdObj != null) {
                                result.setDocumentId(documentIdObj.toString());
                            }
                        }

                        if (contentList != null && i < contentList.size()) {
                            Object contentObj = contentList.get(i);
                            if (contentObj != null) {
                                result.setContent(contentObj.toString());
                            }
                        }

                        if (chunkIndexList != null && i < chunkIndexList.size()) {
                            Object chunkIndexObj = chunkIndexList.get(i);
                            if (chunkIndexObj != null) {
                                if (chunkIndexObj instanceof Long) {
                                    result.setChunkIndex(((Long) chunkIndexObj).intValue());
                                } else if (chunkIndexObj instanceof Integer) {
                                    result.setChunkIndex((Integer) chunkIndexObj);
                                } else if (chunkIndexObj instanceof Number) {
                                    result.setChunkIndex(((Number) chunkIndexObj).intValue());
                                }
                            }
                        }

                        results.add(result);
                        log.debug("成功解析搜索结果 {}: documentId={}, score={}",
                                i, result.getDocumentId(), score);


                    } catch (Exception e) {
                        log.warn("解析索引 {} 的搜索结果失败: {}", i, e.getMessage());
                        continue;
                    }
                }
            }

            log.info("找到 {} 个相似块，阈值 {}", results.size(), threshold);
            return results;

        } catch (Exception e) {
            log.error("在Milvus中搜索相似块失败", e);
            throw new RuntimeException("Failed to search in Milvus", e);
        }
    }

    /**
     * 安全获取字段值的辅助方法
     */
    private Object getFieldValue(SearchResultsWrapper wrapper, String fieldName, int index) {
        try {
            return wrapper.getFieldData(fieldName, index);
        } catch (Exception e1) {
            log.warn("获取索引 {} 的字段 {} 失败: {}", index, fieldName, e1.getMessage());
            return null;
        }
    }
    /**
     * 删除文档相关的所有向量
     */
    public void deleteDocumentChunks(String documentId) {
        try {
            String expression = String.format("document_id == \"%s\"", documentId);

            DeleteParam deleteParam = DeleteParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withExpr(expression)
                    .build();

            milvusClient.delete(deleteParam);
            log.info("删除文档的块: {}", documentId);

        } catch (Exception e) {
            log.error("从Milvus删除文档块失败", e);
            throw new RuntimeException("Failed to delete from Milvus", e);
        }
    }

    @Value("${milvus.codeblock-collection-name:codeblock}")
    private String codeBlockCollectionName;

    /**
     * 插入代码块向量
     */
    public void insertCodeBlocks(List<CodeBlock> codeBlocks) {
        if (codeBlocks.isEmpty()) {
            return;
        }

        try {
            List<String> ids = codeBlocks.stream()
                    .map(CodeBlock::getId)
                    .collect(Collectors.toList());

            List<String> names = codeBlocks.stream()
                    .map(CodeBlock::getName)
                    .collect(Collectors.toList());

            List<String> codes = codeBlocks.stream()
                    .map(CodeBlock::getCode)
                    .collect(Collectors.toList());

            List<String> infos = codeBlocks.stream()
                    .map(CodeBlock::getInfo)
                    .collect(Collectors.toList());

            List<String> languages = codeBlocks.stream()
                    .map(CodeBlock::getLanguage)
                    .collect(Collectors.toList());

            List<String> labelsJson = codeBlocks.stream()
                    .map(block -> String.join(",", block.getLabels()))
                    .collect(Collectors.toList());

            List<List<Float>> vectors = codeBlocks.stream()
                    .map(block -> {
                        List<Float> vector = new ArrayList<>();
                        for (float f : block.getVector()) {
                            vector.add(f);
                        }
                        return vector;
                    })
                    .collect(Collectors.toList());

            List<InsertParam.Field> fields = Arrays.asList(
                    new InsertParam.Field("id", ids),
                    new InsertParam.Field("name", names),
                    new InsertParam.Field("code", codes),
                    new InsertParam.Field("info", infos),
                    new InsertParam.Field("language", languages),
                    new InsertParam.Field("labels", labelsJson),
                    new InsertParam.Field("vector", vectors)
            );

            InsertParam insertParam = InsertParam.newBuilder()
                    .withCollectionName(codeBlockCollectionName)
                    .withFields(fields)
                    .build();

            milvusClient.insert(insertParam);
            log.info("向Milvus插入了 {} 个代码块", codeBlocks.size());

        } catch (Exception e) {
            log.error("向Milvus插入代码块失败", e);
            throw new RuntimeException("Failed to insert code blocks into Milvus", e);
        }
    }



    /**
     * 搜索相似代码块
     */
    /**
     * 在Milvus中搜索相似代码块
     */
    public List<CodeSearchResponse.CodeSearchResult> searchSimilarCodeBlocks(
            float[] queryVector, String language, List<String> labels, int topK, float threshold) {
        try {
            List<String> outputFields = Arrays.asList("id", "name", "code", "info", "language", "labels");

            // 构建过滤表达式
            StringBuilder filterExpr = new StringBuilder();
            if (language != null && !language.isEmpty()) {
                filterExpr.append("language == \"").append(language).append("\"");
            }
            if (labels != null && !labels.isEmpty()) {
                if (filterExpr.length() > 0) {
                    filterExpr.append(" && ");
                }
                List<String> labelConditions = labels.stream()
                        .map(label -> "labels like \"%" + label + "%\"")
                        .collect(Collectors.toList());
                filterExpr.append("(").append(String.join(" || ", labelConditions)).append(")");
            }

            List<Float> vectorList = new ArrayList<>(queryVector.length);
            for (float f : queryVector) {
                vectorList.add(f);
            }
            List<List<Float>> searchVectors = Collections.singletonList(vectorList);

            SearchParam.Builder searchParamBuilder = SearchParam.newBuilder()
                    .withCollectionName(codeBlockCollectionName)
                    .withMetricType(MetricType.COSINE)
                    .withOutFields(outputFields)
                    .withTopK(topK)
                    .withVectors(searchVectors)
                    .withVectorFieldName("vector")
                    .withParams("{\"nprobe\":10}")
                    .withConsistencyLevel(ConsistencyLevelEnum.STRONG);

            if (filterExpr.length() > 0) {
                searchParamBuilder.withExpr(filterExpr.toString());
            }

            SearchParam searchParam = searchParamBuilder.build();
            SearchResults searchResults = milvusClient.search(searchParam).getData();
            SearchResultsWrapper wrapper = new SearchResultsWrapper(searchResults.getResults());

            List<CodeSearchResponse.CodeSearchResult> results = new ArrayList<>();

            if (wrapper.getRowRecords().isEmpty()) {
                log.info("未找到代码块搜索结果");
                return results;
            }

            // 获取分数信息
            List<SearchResultsWrapper.IDScore> idScores = wrapper.getIDScore(0);

            // 获取各字段数据
            List<?> idList = wrapper.getFieldData("id", 0);           // 🔥 通过字段获取ID
            List<?> nameList = wrapper.getFieldData("name", 0);
            List<?> codeList = wrapper.getFieldData("code", 0);
            List<?> infoList = wrapper.getFieldData("info", 0);
            List<?> languageList = wrapper.getFieldData("language", 0);
            List<?> labelsList = wrapper.getFieldData("labels", 0);

            for (int i = 0; i < idScores.size(); i++) {
                SearchResultsWrapper.IDScore idScore = idScores.get(i);
                float score = idScore.getScore();

                if (score >= threshold) {
                    CodeSearchResponse.CodeSearchResult result = new CodeSearchResponse.CodeSearchResult();
                    result.setScore(score);

                    try {
                        // 🔥 正确的ID获取方式
                        if (idList != null && i < idList.size()) {
                            Object idObj = idList.get(i);
                            if (idObj != null) {
                                result.setId(idObj.toString());  // 设置ID
                            }
                        }

                        if (nameList != null && i < nameList.size()) {
                            Object nameObj = nameList.get(i);
                            if (nameObj != null) {
                                result.setName(nameObj.toString());
                            }
                        }

                        if (codeList != null && i < codeList.size()) {
                            Object codeObj = codeList.get(i);
                            if (codeObj != null) {
                                result.setCode(codeObj.toString());
                            }
                        }

                        if (infoList != null && i < infoList.size()) {
                            Object infoObj = infoList.get(i);
                            if (infoObj != null) {
                                result.setInfo(infoObj.toString());
                            }
                        }

                        if (languageList != null && i < languageList.size()) {
                            Object languageObj = languageList.get(i);
                            if (languageObj != null) {
                                result.setLanguage(languageObj.toString());
                            }
                        }

                        if (labelsList != null && i < labelsList.size()) {
                            Object labelsObj = labelsList.get(i);
                            if (labelsObj != null) {
                                String labelsStr = labelsObj.toString();
                                if (!labelsStr.isEmpty()) {
                                    result.setLabels(Arrays.asList(labelsStr.split(",")));
                                } else {
                                    result.setLabels(new ArrayList<>());
                                }
                            }
                        }

                        results.add(result);
                        log.debug("解析代码块搜索结果 {}: id={}, score={}",
                                i, result.getId(), score);

                    } catch (Exception e) {
                        log.warn("解析代码块搜索结果失败: {}", e.getMessage());
                        continue;
                    }
                }
            }

            log.info("找到 {} 个相似代码块，阈值 {}", results.size(), threshold);
            return results;

        } catch (Exception e) {
            log.error("在Milvus中搜索相似代码块失败", e);
            throw new RuntimeException("Failed to search code blocks in Milvus", e);
        }
    }


    /**
     * 删除代码块
     */
    public void deleteCodeBlock(String codeBlockId) {
        try {
            String expression = String.format("id == \"%s\"", codeBlockId);

            DeleteParam deleteParam = DeleteParam.newBuilder()
                    .withCollectionName(codeBlockCollectionName)
                    .withExpr(expression)
                    .build();

            milvusClient.delete(deleteParam);
            log.info("删除代码块: {}", codeBlockId);

        } catch (Exception e) {
            log.error("从Milvus删除代码块失败", e);
            throw new RuntimeException("Failed to delete code block from Milvus", e);
        }
    }
}